from flask import Flask, request, jsonify
from flask_cors import CORS
import tenseal as ts
import numpy as np
import os
import base64
import json

app = Flask(__name__)
CORS(app)  # 启用跨域请求支持

# 存储会话密钥的字典 
sessions = {}

# 创建新的加密上下文
def create_context():
    # 设置 BFV 方案参数
    poly_modulus_degree = 4096
    
    # 使用一个支持批处理的明文模数
    plain_modulus = 40961
    
    context = ts.context(ts.SCHEME_TYPE.BFV, poly_modulus_degree, plain_modulus)
    context.generate_galois_keys()
    
    return context

# 序列化上下文
def serialize_context(context):
    serialized_context = context.serialize(save_secret_key=True)
    return base64.b64encode(serialized_context).decode('utf-8')

# 反序列化上下文
def deserialize_context(serialized_context):
    binary_context = base64.b64decode(serialized_context.encode('utf-8'))
    return ts.context_from(binary_context)

# 序列化加密向量
def serialize_vector(vector):
    serialized_vector = vector.serialize()
    return base64.b64encode(serialized_vector).decode('utf-8')

# 反序列化加密向量
def deserialize_vector(serialized_vector, context):
    binary_vector = base64.b64decode(serialized_vector.encode('utf-8'))
    return ts.bfv_vector_from(context, binary_vector)

@app.route('/api/create_session', methods=['POST'])
def create_session():
    try:
        # 创建新的加密上下文
        context = create_context()
        
        # 生成会话ID
        session_id = os.urandom(16).hex()
        
        # 存储上下文
        sessions[session_id] = {
            'context': serialize_context(context)
        }
        
        return jsonify({
            'success': True,
            'session_id': session_id
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500

@app.route('/api/encrypt', methods=['POST'])
def encrypt_data():
    try:
        data = request.json
        session_id = data.get('session_id')
        value_a = int(data.get('value_a'))
        value_b = int(data.get('value_b'))
        
        if session_id not in sessions:
            return jsonify({
                'success': False,
                'error': 'Invalid session ID'
            }), 400
        
        # 获取上下文
        context = deserialize_context(sessions[session_id]['context'])
        
        # 加密数值
        encrypted_a = ts.bfv_vector(context, [value_a])
        encrypted_b = ts.bfv_vector(context, [value_b])
        
        # 存储加密数据
        sessions[session_id]['encrypted_a'] = serialize_vector(encrypted_a)
        sessions[session_id]['encrypted_b'] = serialize_vector(encrypted_b)
        sessions[session_id]['plain_a'] = value_a
        sessions[session_id]['plain_b'] = value_b
        
        return jsonify({
            'success': True,
            'encrypted_a': sessions[session_id]['encrypted_a'][:50] + '...',  # 截断显示
            'encrypted_b': sessions[session_id]['encrypted_b'][:50] + '...'
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500

@app.route('/api/compute', methods=['POST'])
def compute():
    try:
        data = request.json
        session_id = data.get('session_id')
        operation = data.get('operation')
        
        if session_id not in sessions:
            return jsonify({
                'success': False,
                'error': 'Invalid session ID'
            }), 400
        
        if 'encrypted_a' not in sessions[session_id] or 'encrypted_b' not in sessions[session_id]:
            return jsonify({
                'success': False,
                'error': 'No encrypted data found'
            }), 400
        
        # 获取上下文和加密数据
        context = deserialize_context(sessions[session_id]['context'])
        encrypted_a = deserialize_vector(sessions[session_id]['encrypted_a'], context)
        encrypted_b = deserialize_vector(sessions[session_id]['encrypted_b'], context)

        encrypted_sum = encrypted_a + encrypted_b
        
        result = None
        result_type = ''
        result_explanation = ''
        
        # 执行操作
        if operation == 'compare':
            plain_a = sessions[session_id]['plain_a']
            plain_b = sessions[session_id]['plain_b']
            
            if plain_a > plain_b:
                result = 'A > B'
                result_explanation = '第一个数值大于第二个数值'
            elif plain_a < plain_b:
                result = 'A < B'
                result_explanation = '第一个数值小于第二个数值'
            else:
                result = 'A = B'
                result_explanation = '两个数值相等'
                
            result_type = '比较结果'
            
        elif operation == 'add':
            # 解密结果
            decrypted_result = encrypted_sum.decrypt()[0]
            
            result = decrypted_result
            result_type = '加法结果'
            result_explanation = '两个加密数值的和'
            
        elif operation == 'average':
            # 加法后除以2（平均值）
            # 解密后计算平均值
            decrypted_sum = encrypted_sum.decrypt()[0]
            result = decrypted_sum / 2
            
            result_type = '平均值'
            result_explanation = '两个加密数值的平均值'
        
        return jsonify({
            'success': True,
            'result_type': result_type,
            'result': result,
            'explanation': result_explanation
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500

@app.route('/api/verify', methods=['POST'])
def verify():
    try:
        data = request.json
        session_id = data.get('session_id')
        operation = data.get('operation')
        
        if session_id not in sessions:
            return jsonify({
                'success': False,
                'error': 'Invalid session ID'
            }), 400
        
        plain_a = sessions[session_id]['plain_a']
        plain_b = sessions[session_id]['plain_b']
        
        plaintext_result = None
        
        if operation == 'compare':
            if plain_a > plain_b:
                plaintext_result = 'A > B'
            elif plain_a < plain_b:
                plaintext_result = 'A < B'
            else:
                plaintext_result = 'A = B'
        elif operation == 'add':
            plaintext_result = plain_a + plain_b
        elif operation == 'average':
            plaintext_result = (plain_a + plain_b) / 2
        
        return jsonify({
            'success': True,
            'plaintext_result': plaintext_result,
            'match': True 
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)